package lib

import (
	"algorithms/ds"
)

func init() {
	Probs = append(Probs, Problem{
		Num:         314,
		Discription: "二叉树的垂直遍历（从上到下，逐列）",
		Level:       2,
		Labels: map[string]int{
			"二叉树": 1,
			"哈希表": 1,
		},
	})
}

type RowNode struct {
	Node *ds.TreeNode
	Row  int
}

// 层次遍历，往左走列号-1，往右走列号+1，用哈希表存每一列
func VerticalOrder(root *ds.TreeNode) [][]int {
	if root == nil {
		return [][]int{}
	}

	leftest := 0
	rightest := 0
	rowMap := make(map[int][]int, 0)
	q := []*RowNode{{
		Node: root,
		Row:  0,
	}}
	for len(q) > 0 {
		node := q[0].Node
		row := q[0].Row
		q = q[1:]
		leftest = min(leftest, row)
		rightest = max(rightest, row)
		rowMap[row] = append(rowMap[row], node.Val)
		if node.Left != nil {
			q = append(q, &RowNode{
				Node: node.Left,
				Row:  row - 1,
			})
		}
		if node.Right != nil {
			q = append(q, &RowNode{
				Node: node.Right,
				Row:  row + 1,
			})
		}
	}

	res := make([][]int, 0)
	for i := leftest; i <= rightest; i++ {
		res = append(res, rowMap[i])
	}

	return res
}
